/*
 * Copyright 2023-2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.ai.anthropic;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.springframework.ai.anthropic.api.AnthropicApi;
import org.springframework.ai.anthropic.api.AnthropicApi.ChatCompletionResponse;
import org.springframework.ai.anthropic.api.AnthropicApi.ContentBlock;
import org.springframework.ai.anthropic.api.AnthropicApi.FileMetadata;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.util.Assert;

/**
 * Helper utilities for working with Claude Skills responses and files. Provides methods
 * to extract file IDs, container IDs, and download files generated by Skills.
 *
 * @author Soby Chacko
 * @since 2.0.0
 */
public final class SkillsResponseHelper {

	private SkillsResponseHelper() {
		// Utility class, no instantiation
	}

	/**
	 * Extract all file IDs from a chat response. Searches through all content blocks in
	 * the response, including those in the underlying AnthropicApi response metadata.
	 * @param response The chat response to search
	 * @return List of file IDs found in the response (empty list if none found)
	 */
	public static List<String> extractFileIds(ChatResponse response) {
		if (response == null) {
			return List.of();
		}

		List<String> fileIds = new ArrayList<>();

		// Try to get the underlying Anthropic response from ChatResponse metadata
		if (response.getMetadata() != null) {
			Object anthropicResponse = response.getMetadata().get("anthropic-response");
			if (anthropicResponse instanceof ChatCompletionResponse chatCompletionResponse) {
				fileIds.addAll(extractFileIdsFromContentBlocks(chatCompletionResponse.content()));
			}
		}

		return fileIds;
	}

	/**
	 * Extract container ID from a chat response for multi-turn conversation reuse.
	 * @param response The chat response
	 * @return Container ID if present, null otherwise
	 */
	public static String extractContainerId(ChatResponse response) {
		if (response == null) {
			return null;
		}

		// Try to get container from ChatResponse metadata
		if (response.getMetadata() != null) {
			Object anthropicResponse = response.getMetadata().get("anthropic-response");
			if (anthropicResponse instanceof ChatCompletionResponse chatCompletionResponse) {
				if (chatCompletionResponse.container() != null) {
					return chatCompletionResponse.container().id();
				}
			}
		}

		return null;
	}

	/**
	 * Download all files from a Skills response to a target directory.
	 *
	 * <p>
	 * <b>Note:</b> Existing files with the same name will be overwritten. Check for file
	 * existence before calling if overwrite protection is needed.
	 * @param response The chat response containing file IDs
	 * @param api The Anthropic API client to use for downloading
	 * @param targetDir Directory to save files (must exist)
	 * @return List of paths to saved files
	 * @throws IOException if file download or saving fails
	 */
	public static List<Path> downloadAllFiles(ChatResponse response, AnthropicApi api, Path targetDir)
			throws IOException {
		Assert.notNull(response, "Response cannot be null");
		Assert.notNull(api, "AnthropicApi cannot be null");
		Assert.notNull(targetDir, "Target directory cannot be null");
		Assert.isTrue(Files.isDirectory(targetDir), "Target path must be a directory");

		List<String> fileIds = extractFileIds(response);
		List<Path> savedPaths = new ArrayList<>();

		for (String fileId : fileIds) {
			// Get metadata for filename
			FileMetadata metadata = api.getFileMetadata(fileId);

			// Download file
			byte[] content = api.downloadFile(fileId);

			// Save to target directory
			Path filePath = targetDir.resolve(metadata.filename());
			Files.write(filePath, content);
			savedPaths.add(filePath);
		}

		return savedPaths;
	}

	/**
	 * Extract file IDs from a list of content blocks. Searches both direct file blocks
	 * and nested content structures (for Skills tool results).
	 * @param contentBlocks List of content blocks to search
	 * @return List of file IDs found
	 */
	private static List<String> extractFileIdsFromContentBlocks(List<ContentBlock> contentBlocks) {
		if (contentBlocks == null || contentBlocks.isEmpty()) {
			return List.of();
		}

		List<String> fileIds = new ArrayList<>();

		for (ContentBlock block : contentBlocks) {
			// Check direct fileId field (top-level file blocks)
			if (block.type() == ContentBlock.Type.FILE && block.fileId() != null) {
				fileIds.add(block.fileId());
			}

			// Check nested content field (Skills tool results with complex JSON
			// structures)
			if (block.content() != null) {
				fileIds.addAll(extractFileIdsFromObject(block.content()));
			}
		}

		return fileIds;
	}

	/**
	 * Recursively extract file IDs from any object structure. Handles nested Maps and
	 * Lists that may contain file_id fields deep in the structure (e.g., Skills
	 * bash_code_execution_tool_result responses).
	 * @param obj The object to search (can be Map, List, String, or other types)
	 * @return List of file IDs found in the object structure
	 */
	private static List<String> extractFileIdsFromObject(Object obj) {
		List<String> fileIds = new ArrayList<>();

		if (obj instanceof Map<?, ?> map) {
			// Check if this map has a file_id key
			if (map.containsKey("file_id") && map.get("file_id") instanceof String fileId) {
				fileIds.add(fileId);
			}
			// Recursively search all values in the map
			for (Object value : map.values()) {
				fileIds.addAll(extractFileIdsFromObject(value));
			}
		}
		else if (obj instanceof List<?> list) {
			// Recursively search all list items
			for (Object item : list) {
				fileIds.addAll(extractFileIdsFromObject(item));
			}
		}
		// For String, Number, etc., there are no file_ids to extract

		return fileIds;
	}

}
